import logging

import psycopg2
from sqlalchemy import text

from .connection import get_engine
from ..config.settings import DATABASE_HOST, DATABASE_PORT, DATABASE_USER, DATABASE_PASSWORD, \
    DATABASE_NAME
from ..models.hydrometric import Base, TIMESCALE_SETUP_SQL

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def create_database_if_not_exists():
    """如果数据库不存在则创建"""
    try:
        # 连接到postgres数据库来创建目标数据库
        conn = psycopg2.connect(
            host=DATABASE_HOST,
            port=DATABASE_PORT,
            user=DATABASE_USER,
            password=DATABASE_PASSWORD,
            database='postgres'  # 连接到默认数据库
        )
        conn.autocommit = True
        cursor = conn.cursor()

        # 检查数据库是否存在
        cursor.execute(f"SELECT 1 FROM pg_database WHERE datname = '{DATABASE_NAME}'")
        exists = cursor.fetchone()

        if not exists:
            cursor.execute(f"CREATE DATABASE {DATABASE_NAME}")
            logger.info(f"数据库 '{DATABASE_NAME}' 创建成功")
        else:
            logger.info(f"数据库 '{DATABASE_NAME}' 已存在")

        cursor.close()
        conn.close()

    except Exception as e:
        logger.error(f"创建数据库时出错: {e}")
        raise


def setup_timescaledb():
    """设置TimescaleDB扩展和hypertable"""
    engine = get_engine()
    try:
        with engine.connect() as conn:
            # 执行TimescaleDB设置SQL
            for statement in TIMESCALE_SETUP_SQL.split(';'):
                statement = statement.strip()
                if statement:
                    try:
                        conn.execute(text(statement))
                        conn.commit()
                    except Exception as e:
                        logger.warning(f"执行SQL语句时出现警告: {statement[:50]}... - {e}")
                        # 继续执行其他语句
                        pass

        logger.info("TimescaleDB配置完成")

    except Exception as e:
        logger.error(f"设置TimescaleDB时出错: {e}")
        raise


def drop_tables_if_exists():
    """强制删除现有表（如果存在）"""
    engine = get_engine()
    try:
        with engine.connect() as conn:
            # 删除hypertable和普通表
            drop_statements = [
                "DROP TABLE IF EXISTS hydrometric_data_timeseries CASCADE",
                "DROP TABLE IF EXISTS hydrometric_data_regular CASCADE"
            ]
            
            for statement in drop_statements:
                try:
                    conn.execute(text(statement))
                    conn.commit()
                    logger.info(f"执行删除语句: {statement}")
                except Exception as e:
                    logger.warning(f"删除表时出现警告: {statement} - {e}")
                    
        logger.info("现有表删除完成")
    except Exception as e:
        logger.error(f"删除表时出错: {e}")
        raise


def create_tables():
    """创建所有数据表"""
    engine = get_engine()
    try:
        Base.metadata.create_all(bind=engine)
        logger.info("数据表创建成功")
    except Exception as e:
        logger.error(f"创建数据表时出错: {e}")
        raise


def test_connection():
    """测试数据库连接"""
    engine = get_engine()
    try:
        with engine.connect() as conn:
            result = conn.execute(text("SELECT version()"))
            version = result.fetchone()[0]
            logger.info(f"数据库连接成功: {version}")

            # 检查TimescaleDB扩展
            result = conn.execute(text("SELECT * FROM pg_extension WHERE extname = 'timescaledb'"))
            if result.fetchone():
                logger.info("TimescaleDB扩展已启用")
            else:
                logger.warning("TimescaleDB扩展未启用")

    except Exception as e:
        logger.error(f"数据库连接测试失败: {e}")
        raise


def initialize_database():
    """完整的数据库初始化流程"""
    logger.info("开始数据库初始化...")

    try:
        # 1. 创建数据库（如果不存在）
        create_database_if_not_exists()

        # 2. 测试连接
        test_connection()

        # 3. 强制删除现有表
        drop_tables_if_exists()

        # 4. 创建数据表
        create_tables()

        # 5. 设置TimescaleDB
        setup_timescaledb()

        logger.info("数据库初始化完成！")

    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        raise


if __name__ == "__main__":
    initialize_database()
